import json
import requests
from fastapi import FastAPI, HTTPException, Body
from pydantic import BaseModel, Field
from typing import Dict, Any

from elasticsearch import Elasticsearch

# Import configurations and the prompt template
import config
from prompt_template import PROMPT_TEMPLATE

# --- FastAPI App Initialization ---
app = FastAPI(
    title="Natural Language to Elasticsearch API",
    description="An API to convert natural language to ES queries and execute them.",
    version="1.0.0",
)

# --- Elasticsearch Connection ---
try:
    es_client = Elasticsearch(
        hosts=[{'host': config.ES_HOST, 'port': config.ES_PORT}],
        http_auth=(config.ES_USER, config.ES_PASSWORD),
        use_ssl=config.ES_USE_SSL,
        verify_certs=config.ES_VERIFY_CERTS,
        timeout=30
    )
    # Check connection
    if not es_client.ping():
        raise ConnectionError("Could not connect to Elasticsearch.")
    print("Successfully connected to Elasticsearch.")
except Exception as e:
    print(f"Failed to connect to Elasticsearch: {e}")
    es_client = None

# --- API Request and Response Models ---
class GenerateQueryRequest(BaseModel):
    user_query: str = Field(..., description="The natural language query from the user.")

class GenerateQueryResponse(BaseModel):
    es_query: Dict[str, Any] = Field(..., description="The generated Elasticsearch DSL query.")

class ExecuteQueryRequest(BaseModel):
    index: str = Field(..., description="The target Elasticsearch index (e.g., 'logstash-spat-*').")
    query: Dict[str, Any] = Field(..., description="A valid Elasticsearch DSL query.")


# --- API Endpoints ---

@app.post("/generate-es-query", response_model=GenerateQueryResponse)
def generate_es_query(request: GenerateQueryRequest):
    """
    Receives a natural language query, formats a prompt with context and examples,
    and calls the DeepSeek API to generate an Elasticsearch DSL query.
    """
    print(f"Received user query: {request.user_query}")
    
    # Format the prompt
    final_prompt = PROMPT_TEMPLATE.format(user_query=request.user_query)

    headers = {
        "Authorization": f"Bearer {config.DEEPSEEK_API_KEY}",
        "Content-Type": "application/json",
    }
    
    payload = {
        "model": "deepseek-chat",
        "messages": [
            {"role": "system", "content": "You are an expert assistant for Elasticsearch."},
            {"role": "user", "content": final_prompt},
        ],
        "temperature": 0.1, # Lower temperature for more deterministic output
        "max_tokens": 2048
    }

    try:
        print("Sending request to DeepSeek API...")
        response = requests.post(config.DEEPSEEK_API_URL, headers=headers, json=payload, timeout=60)
        response.raise_for_status()  # Raise an exception for bad status codes (4xx or 5xx)
        
        llm_output = response.json()["choices"][0]["message"]["content"]
        print(f"Received response from DeepSeek: {llm_output}")

        # Clean up the response to get only the JSON part
        json_str = llm_output.strip().replace("```json", "").replace("```", "").strip()

        # Parse the JSON string into a Python dictionary
        es_query_dict = json.loads(json_str)
        
        return {"es_query": es_query_dict}

    except requests.exceptions.RequestException as e:
        print(f"Error calling DeepSeek API: {e}")
        raise HTTPException(status_code=500, detail=f"Failed to communicate with DeepSeek API: {e}")
    except (json.JSONDecodeError, KeyError, IndexError) as e:
        print(f"Error parsing LLM response: {e}")
        raise HTTPException(status_code=500, detail=f"Failed to parse the response from the language model. Raw output: {llm_output}")


@app.post("/execute-es-query")
def execute_es_query(request: ExecuteQueryRequest):
    """
    Receives an Elasticsearch index and DSL query, and executes it against the database.
    """
    if es_client is None:
        raise HTTPException(status_code=503, detail="Elasticsearch client is not available.")

    print(f"Executing query on index '{request.index}'")
    
    try:
        search_results = es_client.search(index=request.index, body=request.query)
        return search_results
    except Exception as e:
        print(f"Error executing Elasticsearch query: {e}")
        # Pass the raw ES error to the client for better debugging
        raise HTTPException(status_code=400, detail=str(e)) 